import uuid
from typing import List, Dict, Optional

import pandas as pd
from pymilvus import DataType, MilvusClient, model, FieldSchema, CollectionSchema

# 移除 VannaBase 的继承
# from vanna.base import VannaBase

# Setting the URI as a local file, e.g.`./milvus.db`,
# is the most convenient method, as it automatically utilizes Milvus Lite
# to store all data in this file.
DEFAULT_MILVUS_URI = "./milvus.db"
MAX_LIMIT_SIZE = 10_000


class SqlTemplateStore:
    """
    SQL Template store implementation using Milvus for managing SQL templates with embeddings.

    Args:
        - config (dict, optional): Dictionary of configuration options. Defaults to `None`.
            - milvus_client: A `pymilvus.MilvusClient` instance.
            - embedding_function: A `milvus_model.base.BaseEmbeddingFunction` instance.
            - dimension: Embedding dimension. Defaults to 768.
    """

    def __init__(self, config=None):
        if config is None:
            config = {}

        if "milvus_client" in config:
            self.milvus_client = config["milvus_client"]
        else:
            self.milvus_client = MilvusClient(uri=DEFAULT_MILVUS_URI)

        if "embedding_function" in config:
            self.embedding_function = config.get("embedding_function")
        else:
            self.embedding_function = model.DefaultEmbeddingFunction()

        # Get embedding dimension
        if "dimension" in config:
            self._embedding_dim = config["dimension"]
        else:
            self._embedding_dim = self.embedding_function.encode_documents(["test"])[0].shape[0]

        self.collection_name = "sql_template"
        self.n_results = config.get("n_results", 10)

        # Create the sql_template collection
        self._create_sql_template_collection()

    def _create_sql_template_collection(self):
        """Create the sql_template collection if it doesn't exist."""
        if not self.milvus_client.has_collection(collection_name=self.collection_name):
            # Create schema
            schema = MilvusClient.create_schema(
                auto_id=False,
                enable_dynamic_field=False,
            )

            # Add fields
            schema.add_field(field_name="id", datatype=DataType.VARCHAR, max_length=100, is_primary=True)
            schema.add_field(field_name="description", datatype=DataType.VARCHAR, max_length=2000)
            schema.add_field(field_name="sql_template", datatype=DataType.VARCHAR, max_length=4000)
            schema.add_field(field_name="parameters", datatype=DataType.VARCHAR, max_length=4000)
            schema.add_field(field_name="plotly_code", datatype=DataType.VARCHAR, max_length=4000)
            schema.add_field(field_name="embedding", datatype=DataType.FLOAT_VECTOR, dim=self._embedding_dim)

            # Create index
            index_params = self.milvus_client.prepare_index_params()
            index_params.add_index(
                field_name="embedding",
                index_name="embedding_index",
                index_type="AUTOINDEX",
                metric_type="COSINE",
            )

            # Create collection
            self.milvus_client.create_collection(
                collection_name=self.collection_name,
                schema=schema,
                index_params=index_params,
                consistency_level="Strong"
            )

    def generate_embedding(self, text: str) -> List[float]:
        """Generate embedding for given text."""
        return self.embedding_function.encode_documents([text])[0].tolist()

    def add_sql_template(self,
                         id: str,
                         description: str,
                         sql_template: str,
                         parameters: str = "",
                         plotly_code: str = "",
                         **kwargs) -> str:
        """
        Add a new SQL template to the collection.

        Args:
            id: Unique identifier for the template
            description: Description of the SQL template
            sql_template: The SQL template string
            parameters: Parameters for the template (JSON string or description)
            plotly_code: Plotly visualization code

        Returns:
            str: The ID of the inserted template
        """
        if not id or not description or not sql_template:
            raise ValueError("id, description, and sql_template cannot be empty")

        # Generate embedding based on description
        embedding = self.embedding_function.encode_documents([description])[0]

        # Insert data
        self.milvus_client.insert(
            collection_name=self.collection_name,
            data={
                "id": id,
                "description": description,
                "sql_template": sql_template,
                "parameters": parameters,
                "plotly_code": plotly_code,
                "embedding": embedding
            }
        )
        return id

    def get_sql_template_by_id(self, id: str) -> Optional[Dict]:
        """
        Get a SQL template by its ID.

        Args:
            id: The ID of the template to retrieve

        Returns:
            Dict or None: The template data if found, None otherwise
        """
        try:
            results = self.milvus_client.query(
                collection_name=self.collection_name,
                filter=f'id == "{id}"',
                output_fields=["*"]
            )

            if results:
                return results[0]
            return None
        except Exception as e:
            print(f"Error retrieving template {id}: {e}")
            return None

    def search_similar_templates(self, query: str, limit: Optional[int] = None, score_threshold: Optional[float] = 0.85) -> List[Dict]:
        """
        Search for similar SQL templates based on description.

        Args:
            query: Query string to search for
            limit: Maximum number of results to return
            score_threshold: Minimum score threshold for results
        Returns:
            List[Dict]: List of similar templates
        """
        if limit is None:
            limit = self.n_results

        search_params = {
            "metric_type": "COSINE",
            "params": {"nprobe": 128},
        }

        # Generate embedding for query
        query_embedding = self.embedding_function.encode_queries([query])

        # Search
        results = self.milvus_client.search(
            collection_name=self.collection_name,
            anns_field="embedding",
            data=query_embedding,
            limit=limit,
            output_fields=["id", "description", "sql_template", "parameters", "plotly_code"],
            search_params=search_params
        )
        print(f"search_similar_templates:{results}")
        # Format results
        templates = []
        for result in results[0]:
            similarity_score = result["distance"]
            # 应用阈值筛选
            if score_threshold is not None:
                if search_params["metric_type"] == "COSINE":
                    # 余弦相似度：值越大越相似
                    if similarity_score < score_threshold:
                        continue
                else:
                    # L2距离：值越小越相似
                    if similarity_score > score_threshold:
                        continue
            template = {
                "id": result["entity"]["id"],
                "description": result["entity"]["description"],
                "sql_template": result["entity"]["sql_template"],
                "parameters": result["entity"]["parameters"],
                "plotly_code": result["entity"]["plotly_code"],
            }
            templates.append(template)

        return templates

    def update_sql_template(self,
                            id: str,
                            description: Optional[str] = None,
                            sql_template: Optional[str] = None,
                            parameters: Optional[str] = None,
                            plotly_code: Optional[str] = None,
                            **kwargs) -> bool:
        """
        Update an existing SQL template.

        Args:
            id: ID of the template to update
            description: New description (optional)
            sql_template: New SQL template (optional)
            parameters: New parameters (optional)
            plotly_code: New plotly code (optional)

        Returns:
            bool: True if update was successful, False otherwise
        """
        try:
            # First, get the existing template
            existing = self.get_sql_template_by_id(id)
            if not existing:
                return False

            # Prepare updated data
            updated_data = {
                "id": id,
                "description": description if description is not None else existing["description"],
                "sql_template": sql_template if sql_template is not None else existing["sql_template"],
                "parameters": parameters if parameters is not None else existing["parameters"],
                "plotly_code": plotly_code if plotly_code is not None else existing["plotly_code"]
            }

            # Generate new embedding if description changed
            if description is not None:
                updated_data["embedding"] = self.embedding_function.encode_documents([description])[0]
            else:
                updated_data["embedding"] = existing["embedding"]

            # Delete old record and insert new one (Milvus doesn't support direct update)
            self.milvus_client.delete(collection_name=self.collection_name, ids=[id])
            self.milvus_client.insert(collection_name=self.collection_name, data=updated_data)

            return True
        except Exception as e:
            print(f"Error updating template {id}: {e}")
            return False

    def delete_sql_template(self, id: str) -> bool:
        """
        Delete a SQL template by its ID.

        Args:
            id: ID of the template to delete

        Returns:
            bool: True if deletion was successful, False otherwise
        """
        try:
            self.milvus_client.delete(collection_name=self.collection_name, ids=[id])
            return True
        except Exception as e:
            print(f"Error deleting template {id}: {e}")
            return False

    def list_all_templates(self, limit: int = MAX_LIMIT_SIZE) -> pd.DataFrame:
        """
        Get all SQL templates as a DataFrame.

        Args:
            limit: Maximum number of templates to retrieve

        Returns:
            pd.DataFrame: DataFrame containing all templates
        """
        try:
            results = self.milvus_client.query(
                collection_name=self.collection_name,
                output_fields=["id", "description", "sql_template", "parameters", "plotly_code"],
                limit=limit
            )

            if not results:
                return pd.DataFrame(columns=["id", "description", "sql_template", "parameters", "plotly_code"])

            df = pd.DataFrame({
                "id": [doc["id"] for doc in results],
                "description": [doc["description"] for doc in results],
                "sql_template": [doc["sql_template"] for doc in results],
                "parameters": [doc["parameters"] for doc in results],
                "plotly_code": [doc["plotly_code"] for doc in results]
            })

            return df
        except Exception as e:
            print(f"Error listing templates: {e}")
            return pd.DataFrame()

    def get_collection_stats(self) -> Dict:
        """
        Get statistics about the sql_template collection.

        Returns:
            Dict: Collection statistics
        """
        try:
            # Get collection info
            collection_info = self.milvus_client.describe_collection(collection_name=self.collection_name)

            # Count total entities
            results = self.milvus_client.query(
                collection_name=self.collection_name,
                output_fields=["id"],
                limit=MAX_LIMIT_SIZE
            )

            return {
                "collection_name": self.collection_name,
                "total_entities": len(results),
                "embedding_dimension": self._embedding_dim
            }
        except Exception as e:
            print(f"Error getting collection stats: {e}")
            return {}
